import pandas as pd
import numpy as np
from scipy.stats.mstats import winsorize
import os
from tqdm import tqdm

# 配置参数
input_dir = r"F:\量化投资程序\机器算法准备\综合修改版\股票数据"
summary_file = os.path.join(input_dir, "文件目录汇总.xlsx")
output_dir = r"F:\上市公司综合数据"
output_file = os.path.join(output_dir, "股票年化收益率统计.csv")
start_year = 2000
end_year = 2023


def load_file_list():
    """加载文件目录汇总"""
    try:
        df = pd.read_excel(summary_file)
        # 假设汇总表中有"文件路径"列，包含相对路径
        return [os.path.join(input_dir, f) for f in df['文件名'] if f.endswith('.csv')]
    except Exception as e:
        print(f"文件列表加载失败: {str(e)}")
        return []


def calculate_annualized_return(file_path):
    """增强版年化收益率计算"""
    try:
        # 读取数据
        df = pd.read_csv(file_path, parse_dates=['trade_date'], index_col='trade_date')
        df.sort_index(inplace=True)

        # 数据过滤
        mask = (df.index.year >= start_year) & (df.index.year <= end_year)
        df = df.loc[mask]

        if df.empty:
            raise ValueError("无有效数据")

        # 数据清洗
        df['close_clean'] = winsorize(df['close'], limits=[0.01, 0.01])

        # 计算对数收益率
        df['log_return'] = np.log(df['close_clean'] / df['close_clean'].shift(1))
        df.dropna(subset=['log_return'], inplace=True)

        # 计算年化收益率
        daily_mean = df['log_return'].mean()
        annualized_mean = daily_mean * 240  # 240个交易日

        total_return = np.exp(df['log_return'].sum()) - 1
        years = (df.index[-1] - df.index[0]).days / 365.25
        annualized_geo = (1 + total_return) ** (1 / years) - 1 if years > 0 else 0

        return {
            "股票代码": os.path.basename(file_path).split('.')[0],
            "起始日期": df.index[0].strftime('%Y-%m-%d'),
            "结束日期": df.index[-1].strftime('%Y-%m-%d'),
            "交易日数": len(df),
            "均值法年化": annualized_mean,
            "几何平均法年化": annualized_geo
        }
    except Exception as e:
        print(f"\n{os.path.basename(file_path)} 处理失败: {str(e)[:50]}")
        return None


def batch_process():
    """批量处理主程序"""
    file_list = load_file_list()
    if not file_list:
        return

    results = []
    error_log = []

    # 创建进度条
    pbar = tqdm(file_list, desc="处理进度", unit="file")

    for file_path in pbar:
        pbar.set_postfix_str(f"正在处理 {os.path.basename(file_path)}")
        try:
            if not os.path.exists(file_path):
                raise FileNotFoundError("文件不存在")

            result = calculate_annualized_return(file_path)
            if result:
                results.append(result)
        except Exception as e:
            error_log.append({
                "文件路径": file_path,
                "错误信息": str(e)[:100]
            })

    # 保存结果
    if results:
        result_df = pd.DataFrame(results)
        result_df.sort_values("股票代码", inplace=True)

        # 格式化显示
        result_df['均值法年化'] = result_df['均值法年化'].apply(lambda x: f"{x:.2%}")
        result_df['几何平均法年化'] = result_df['几何平均法年化'].apply(lambda x: f"{x:.2%}")

        result_df.to_csv(output_file, index=False, encoding='utf_8_sig')
        print(f"\n成功处理 {len(results)}/{len(file_list)} 个文件")
        print(f"结果文件: {output_file}")
    else:
        print("\n未成功处理任何文件")

    # 保存错误日志
    if error_log:
        error_df = pd.DataFrame(error_log)
        error_file = os.path.join(output_dir, "处理错误日志.csv")
        error_df.to_csv(error_file, index=False, encoding='utf_8_sig')
        print(f"错误日志: {error_file}")


if __name__ == "__main__":
    # 检查目录有效性
    if not os.path.exists(input_dir):
        print(f"输入目录不存在: {input_dir}")
        exit()

    if not os.path.exists(summary_file):
        print(f"汇总文件不存在: {summary_file}")
        exit()

    os.makedirs(output_dir, exist_ok=True)

    # 执行批处理
    batch_process()
